"""
模型对比数据访问层
处理MongoDB中的对比会话存储和查询
"""

from typing import List, Optional, Dict, Any
from datetime import datetime, timedelta
from motor.motor_asyncio import AsyncIOMotorClient, AsyncIOMotorDatabase
from bson import ObjectId
import logging

from ..models.comparison import (
    ComparisonSession,
    ComparisonStatus,
    ModelResult,
    ModelSelectionEvent
)

logger = logging.getLogger(__name__)


class ComparisonRepository:
    """
    模型对比数据仓库
    管理对比会话和相关数据的持久化
    """
    
    def __init__(self, db: AsyncIOMotorDatabase):
        """
        初始化仓库
        
        Args:
            db: MongoDB数据库实例
        """
        self.db = db
        self.collection = db.model_comparisons
        self.events_collection = db.comparison_events
        
        # 创建索引
        self._create_indexes()
        
    def _create_indexes(self):
        """创建数据库索引以优化查询性能"""
        # 会话索引
        self.collection.create_index("session_id", unique=True)
        self.collection.create_index("user_id")
        self.collection.create_index("created_at", background=True)
        self.collection.create_index([("user_id", 1), ("created_at", -1)])
        self.collection.create_index("status")
        
        # 事件索引
        self.events_collection.create_index("session_id")
        self.events_collection.create_index("user_id")
        self.events_collection.create_index("timestamp")
        self.events_collection.create_index([("user_id", 1), ("timestamp", -1)])
        
    async def create_session(self, session: ComparisonSession) -> str:
        """
        创建新的对比会话
        
        Args:
            session: 对比会话对象
            
        Returns:
            会话ID
        """
        try:
            doc = session.to_mongo()
            result = await self.collection.insert_one(doc)
            logger.info(f"创建对比会话: {session.session_id}")
            return str(result.inserted_id)
        except Exception as e:
            logger.error(f"创建对比会话失败: {e}")
            raise
            
    async def get_session(self, session_id: str) -> Optional[ComparisonSession]:
        """
        获取对比会话
        
        Args:
            session_id: 会话ID
            
        Returns:
            对比会话对象或None
        """
        try:
            doc = await self.collection.find_one({"session_id": session_id})
            if doc:
                return ComparisonSession.from_mongo(doc)
            return None
        except Exception as e:
            logger.error(f"获取对比会话失败: {e}")
            raise
            
    async def update_session(self, 
                           session_id: str,
                           updates: Dict[str, Any]) -> bool:
        """
        更新对比会话
        
        Args:
            session_id: 会话ID
            updates: 更新字段字典
            
        Returns:
            是否更新成功
        """
        try:
            # 自动更新时间戳
            updates["updated_at"] = datetime.utcnow()
            
            result = await self.collection.update_one(
                {"session_id": session_id},
                {"$set": updates}
            )
            
            if result.modified_count > 0:
                logger.info(f"更新对比会话: {session_id}")
                return True
            return False
        except Exception as e:
            logger.error(f"更新对比会话失败: {e}")
            raise
            
    async def add_result(self, 
                        session_id: str,
                        result: ModelResult) -> bool:
        """
        添加模型结果到会话
        
        Args:
            session_id: 会话ID
            result: 模型结果
            
        Returns:
            是否添加成功
        """
        try:
            update_result = await self.collection.update_one(
                {"session_id": session_id},
                {
                    "$push": {"results": result.dict()},
                    "$set": {"updated_at": datetime.utcnow()}
                }
            )
            
            if update_result.modified_count > 0:
                logger.info(f"添加结果到会话 {session_id}: {result.model}")
                return True
            return False
        except Exception as e:
            logger.error(f"添加结果失败: {e}")
            raise
            
    async def update_status(self,
                          session_id: str,
                          status: ComparisonStatus) -> bool:
        """
        更新会话状态
        
        Args:
            session_id: 会话ID
            status: 新状态
            
        Returns:
            是否更新成功
        """
        return await self.update_session(
            session_id,
            {"status": status.value}
        )
        
    async def set_selected_model(self,
                                session_id: str,
                                model_name: str) -> bool:
        """
        设置用户选择的模型
        
        Args:
            session_id: 会话ID
            model_name: 选择的模型名称
            
        Returns:
            是否设置成功
        """
        return await self.update_session(
            session_id,
            {"selected": model_name}
        )
        
    async def save_merged_content(self,
                                 session_id: str,
                                 content: str) -> bool:
        """
        保存合并后的内容
        
        Args:
            session_id: 会话ID
            content: 合并的内容
            
        Returns:
            是否保存成功
        """
        return await self.update_session(
            session_id,
            {"merged_content": content}
        )
        
    async def get_user_history(self,
                              user_id: str,
                              limit: int = 20,
                              offset: int = 0,
                              status_filter: Optional[ComparisonStatus] = None) -> List[ComparisonSession]:
        """
        获取用户的对比历史
        
        Args:
            user_id: 用户ID
            limit: 返回数量限制
            offset: 偏移量
            status_filter: 状态过滤器
            
        Returns:
            对比会话列表
        """
        try:
            query = {"user_id": user_id}
            if status_filter:
                query["status"] = status_filter.value
                
            cursor = self.collection.find(query).sort("created_at", -1).skip(offset).limit(limit)
            
            sessions = []
            async for doc in cursor:
                sessions.append(ComparisonSession.from_mongo(doc))
                
            logger.info(f"获取用户 {user_id} 的历史记录: {len(sessions)} 条")
            return sessions
        except Exception as e:
            logger.error(f"获取用户历史失败: {e}")
            raise
            
    async def get_recent_comparisons(self,
                                    hours: int = 24,
                                    limit: int = 100) -> List[ComparisonSession]:
        """
        获取最近的对比会话
        
        Args:
            hours: 时间范围（小时）
            limit: 返回数量限制
            
        Returns:
            对比会话列表
        """
        try:
            since = datetime.utcnow() - timedelta(hours=hours)
            cursor = self.collection.find(
                {"created_at": {"$gte": since}}
            ).sort("created_at", -1).limit(limit)
            
            sessions = []
            async for doc in cursor:
                sessions.append(ComparisonSession.from_mongo(doc))
                
            return sessions
        except Exception as e:
            logger.error(f"获取最近对比失败: {e}")
            raise
            
    async def record_selection_event(self, event: ModelSelectionEvent) -> str:
        """
        记录模型选择事件（用于A/B测试）
        
        Args:
            event: 选择事件
            
        Returns:
            事件ID
        """
        try:
            doc = event.dict()
            result = await self.events_collection.insert_one(doc)
            logger.info(f"记录选择事件: {event.session_id} -> {event.selected_model}")
            return str(result.inserted_id)
        except Exception as e:
            logger.error(f"记录选择事件失败: {e}")
            raise
            
    async def get_selection_statistics(self,
                                      user_id: Optional[str] = None,
                                      days: int = 7) -> Dict[str, Any]:
        """
        获取模型选择统计（A/B测试数据）
        
        Args:
            user_id: 用户ID（可选）
            days: 统计天数
            
        Returns:
            统计数据字典
        """
        try:
            since = datetime.utcnow() - timedelta(days=days)
            
            # 构建查询
            match_stage = {"timestamp": {"$gte": since}}
            if user_id:
                match_stage["user_id"] = user_id
                
            # 聚合统计
            pipeline = [
                {"$match": match_stage},
                {"$group": {
                    "_id": "$selected_model",
                    "count": {"$sum": 1},
                    "avg_score": {"$avg": "$scores.$selected_model"}
                }},
                {"$sort": {"count": -1}}
            ]
            
            cursor = self.events_collection.aggregate(pipeline)
            
            stats = {
                "model_selections": {},
                "total_selections": 0,
                "period_days": days
            }
            
            async for doc in cursor:
                model = doc["_id"]
                count = doc["count"]
                stats["model_selections"][model] = {
                    "count": count,
                    "avg_score": doc.get("avg_score", 0)
                }
                stats["total_selections"] += count
                
            # 计算选择百分比
            for model, data in stats["model_selections"].items():
                data["percentage"] = (data["count"] / stats["total_selections"] * 100) if stats["total_selections"] > 0 else 0
                
            logger.info(f"获取选择统计: {stats['total_selections']} 次选择")
            return stats
        except Exception as e:
            logger.error(f"获取选择统计失败: {e}")
            raise
            
    async def cleanup_old_sessions(self, days: int = 30) -> int:
        """
        清理旧的对比会话
        
        Args:
            days: 保留天数
            
        Returns:
            删除的文档数量
        """
        try:
            cutoff = datetime.utcnow() - timedelta(days=days)
            result = await self.collection.delete_many(
                {"created_at": {"$lt": cutoff}}
            )
            
            if result.deleted_count > 0:
                logger.info(f"清理了 {result.deleted_count} 个旧会话")
                
            return result.deleted_count
        except Exception as e:
            logger.error(f"清理旧会话失败: {e}")
            raise